# -*- coding: utf-8 -*-
"""
数据库查询优化工具
提供分页、批量操作、查询优化等功能
"""

from typing import List, Dict, Any, Optional, Tuple
from sqlalchemy import and_, or_, desc, asc, func
from sqlalchemy.orm import joinedload, selectinload, subqueryload
from app import db


class QueryOptimizer:
    """查询优化器"""
    
    @staticmethod
    def paginate_query(query, page: int = 1, per_page: int = 20, order_by: str = None, order_desc: bool = True):
        """
        分页查询
        :param query: SQLAlchemy查询对象
        :param page: 页码（从1开始）
        :param per_page: 每页数量
        :param order_by: 排序字段
        :param order_desc: 是否降序
        :return: (数据列表, 总数, 分页信息)
        """
        # 添加排序
        if order_by:
            if order_desc:
                query = query.order_by(desc(order_by))
            else:
                query = query.order_by(asc(order_by))
        
        # 计算总数
        total = query.count()
        
        # 分页
        offset = (page - 1) * per_page
        items = query.offset(offset).limit(per_page).all()
        
        # 分页信息
        pagination = {
            'page': page,
            'per_page': per_page,
            'total': total,
            'pages': (total + per_page - 1) // per_page,
            'has_prev': page > 1,
            'has_next': page * per_page < total,
            'prev_num': page - 1 if page > 1 else None,
            'next_num': page + 1 if page * per_page < total else None
        }
        
        return items, total, pagination
    
    @staticmethod
    def batch_update(model_class, updates: List[Dict[str, Any]], batch_size: int = 100):
        """
        批量更新
        :param model_class: 模型类
        :param updates: 更新数据列表 [{'id': 1, 'field': 'value'}, ...]
        :param batch_size: 批次大小
        :return: 更新成功的数量
        """
        success_count = 0
        
        for i in range(0, len(updates), batch_size):
            batch = updates[i:i + batch_size]
            
            try:
                for item in batch:
                    record_id = item.pop('id')
                    record = model_class.query.get(record_id)
                    if record:
                        for key, value in item.items():
                            if hasattr(record, key):
                                setattr(record, key, value)
                        success_count += 1
                
                db.session.commit()
                
            except Exception as e:
                db.session.rollback()
                raise e
        
        return success_count
    
    @staticmethod
    def batch_insert(model_class, records: List[Dict[str, Any]], batch_size: int = 100):
        """
        批量插入
        :param model_class: 模型类
        :param records: 记录数据列表
        :param batch_size: 批次大小
        :return: 插入成功的数量
        """
        success_count = 0
        
        for i in range(0, len(records), batch_size):
            batch = records[i:i + batch_size]
            
            try:
                for record_data in batch:
                    record = model_class(**record_data)
                    db.session.add(record)
                    success_count += 1
                
                db.session.commit()
                
            except Exception as e:
                db.session.rollback()
                raise e
        
        return success_count
    
    @staticmethod
    def optimize_joins(query, *relationships):
        """
        优化关联查询
        :param query: 查询对象
        :param relationships: 关联关系列表
        :return: 优化后的查询
        """
        for relationship in relationships:
            query = query.options(joinedload(relationship))
        return query
    
    @staticmethod
    def select_in_load(query, *relationships):
        """
        使用selectinload优化关联查询
        :param query: 查询对象
        :param relationships: 关联关系列表
        :return: 优化后的查询
        """
        for relationship in relationships:
            query = query.options(selectinload(relationship))
        return query


class QueryBuilder:
    """查询构建器"""
    
    def __init__(self, model_class):
        self.model_class = model_class
        self.query = model_class.query
    
    def filter_by_dict(self, filters: Dict[str, Any]):
        """
        根据字典过滤
        :param filters: 过滤条件字典
        :return: 查询构建器
        """
        for key, value in filters.items():
            if value is not None:
                if isinstance(value, (list, tuple)):
                    self.query = self.query.filter(getattr(self.model_class, key).in_(value))
                else:
                    self.query = self.query.filter(getattr(self.model_class, key) == value)
        return self
    
    def search(self, field: str, keyword: str, exact: bool = False):
        """
        搜索
        :param field: 搜索字段
        :param keyword: 搜索关键词
        :param exact: 是否精确匹配
        :return: 查询构建器
        """
        if keyword:
            if exact:
                self.query = self.query.filter(getattr(self.model_class, field) == keyword)
            else:
                self.query = self.query.filter(getattr(self.model_class, field).like(f'%{keyword}%'))
        return self
    
    def date_range(self, field: str, start_date: str = None, end_date: str = None):
        """
        日期范围过滤
        :param field: 日期字段
        :param start_date: 开始日期
        :param end_date: 结束日期
        :return: 查询构建器
        """
        if start_date:
            self.query = self.query.filter(getattr(self.model_class, field) >= start_date)
        if end_date:
            self.query = self.query.filter(getattr(self.model_class, field) <= end_date)
        return self
    
    def order_by_field(self, field: str, desc: bool = True):
        """
        排序
        :param field: 排序字段
        :param desc: 是否降序
        :return: 查询构建器
        """
        if desc:
            self.query = self.query.order_by(desc(getattr(self.model_class, field)))
        else:
            self.query = self.query.order_by(asc(getattr(self.model_class, field)))
        return self
    
    def paginate(self, page: int = 1, per_page: int = 20):
        """
        分页
        :param page: 页码
        :param per_page: 每页数量
        :return: 分页结果
        """
        return QueryOptimizer.paginate_query(self.query, page, per_page)
    
    def all(self):
        """获取所有结果"""
        return self.query.all()
    
    def first(self):
        """获取第一个结果"""
        return self.query.first()
    
    def count(self):
        """获取总数"""
        return self.query.count() 